import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.decomposition import TruncatedSVD
from sklearn.metrics.pairwise import cosine_similarity
from PIL import Image
from sklearn.preprocessing import normalize
import random
import sys

# 1. 文本与图像特征生成
def generate_text_features(documents):
    """ 
基于TF-IDF和SVD生成文本特征
:param documents: 文本列表
:return: 文本嵌入特征矩阵
""" 
    tfidf = TfidfVectorizer(max_features=100)
    tfidf_matrix = tfidf.fit_transform(documents)
    svd = TruncatedSVD(n_components=10, random_state=42)
    return normalize(svd.fit_transform(tfidf_matrix))

def generate_image_features(num_images=10, dimensions=10):
    """ 
生成随机图像嵌入特征
:param num_images: 图像数量
:param dimensions: 嵌入向量维度
:return: 图像嵌入特征矩阵
""" 
    np.random.seed(42)
    return normalize(np.random.rand(num_images, dimensions))

# 示例文本和图像
documents = [
    "apple banana orange",
    "cat dog elephant", 
    "car bus train",
    "coffee tea water",
    "python java c++",
    "rose lily tulip",
    "sun moon stars",
    "pen pencil eraser",
    "chair table desk",
    "lamp fan bulb"
]
text_features = generate_text_features(documents)
image_features = generate_image_features(num_images=10, dimensions=10)

# 2. LSH类定义
class LSH:
    """ 
    实现LSH，用于文本和图像检索
    """ 
    def __init__(self, dimensions, num_hashes, num_tables):
        self.dimensions = dimensions
        self.num_hashes = num_hashes
        self.num_tables = num_tables
        self.hash_planes = [np.random.randn(num_hashes, dimensions) for _ in range(num_tables)]
        self.tables = [{} for _ in range(num_tables)]  # 修正：使用列表推导式创建字典列表

    def hash_function(self, vector, planes):
        projections = np.dot(planes, vector)
        return ''.join(['1' if p > 0 else '0' for p in projections])

    def insert(self, vectors):
        for table_id, planes in enumerate(self.hash_planes):
            for idx, vector in enumerate(vectors):
                hash_value = self.hash_function(vector, planes)
                if hash_value not in self.tables[table_id]:
                    self.tables[table_id][hash_value] = []
                self.tables[table_id][hash_value].append(idx)

    def query(self, query_vector):
        candidates = set()
        for table_id, planes in enumerate(self.hash_planes):
            hash_value = self.hash_function(query_vector, planes)
            if hash_value in self.tables[table_id]:
                candidates.update(self.tables[table_id][hash_value])
        return list(candidates)

# 3. 初始化LSH并插入特征
num_hashes = 5
num_tables = 3
dimensions = 10

# 创建LSH实例
lsh = LSH(dimensions=dimensions, num_hashes=num_hashes, num_tables=num_tables)

# 合并文本和图像特征进行插入
all_features = np.vstack([text_features, image_features])
lsh.insert(all_features)

# 示例查询
query_vector = text_features[0]  # 使用第一个文本特征作为查询
candidates = lsh.query(query_vector)
print("候选索引:", candidates)